
#include <stdio.h>
#include <memory.h>
#include <math.h>

#include <libxml/xmlmemory.h>
#include <libxml/parser.h>

#include <archive.h>
#include <archive_entry.h>

#include <set>
#include <unordered_map>

#include "book.h"
#include "words.h"


struct FB2Parser {

    Book book;

    bool in_title_info;
    bool in_author;
    bool in_first_name;
    bool in_middle_name;
    bool in_last_name;
    bool in_book_title;
    bool in_lang;
    bool in_body;
    bool in_title;
    bool in_subtitle;
    bool in_p;
    bool in_v;
    bool in_text_author;
    bool in_a;

    const std::set<std::string>& langs;
    bool ok;

    FB2Parser(const std::set<std::string>& l) : 
        in_title_info(false),
        in_author(false),
        in_first_name(false),
        in_middle_name(false),
        in_last_name(false),
        in_book_title(false),
        in_lang(false),
        in_body(false),
        in_title(false),
        in_subtitle(false),
        in_p(false),
        in_v(false),
        in_text_author(false),
        in_a(false),
        langs(l),
        ok(true)
        {}


    static void startElementNs(void* ctx, 
                              const xmlChar* localname, 
                              const xmlChar* prefix, 
                              const xmlChar* URI, 
                              int nb_namespaces, 
                              const xmlChar** namespaces, 
                              int nb_attributes, 
                              int nb_defaulted, 
                              const xmlChar** attributes) {

        FB2Parser& p = *(static_cast<FB2Parser*>(ctx));

        std::string name((char*)localname);

        if (name == "title-info") {
            p.in_title_info = true;

        } else if (name == "author") {
            p.in_author = true;

            if (p.in_title_info) {
                p.book.author.push_back(std::vector<std::string>(3));
            }

        } else if (name == "first-name") {
            p.in_first_name = true;

        } else if (name == "middle-name") {
            p.in_middle_name = true;

        } else if (name == "last-name") {
            p.in_last_name = true;

        } else if (name == "book-title") {
            p.in_book_title = true;

        } else if (name == "lang") {
            p.in_lang = true;

        } else if (name == "body") {
            p.in_body = true;

        } else if (name == "title") {
            p.in_title = true;

        } else if (name == "subtitle") {
            p.in_subtitle = true;

            p.book.text.push_back(std::make_pair(Book::LABEL, ""));

        } else if (name == "p") {
            p.in_p = true;

            p.book.text.push_back(std::make_pair(p.in_title ? Book::TITLE : Book::PLAIN, ""));

        } else if (name == "v") {
            p.in_v = true;

            p.book.text.push_back(std::make_pair(Book::POEM, ""));

        } else if (name == "text-author") {
            p.in_text_author = true;

            p.book.text.push_back(std::make_pair(Book::LABEL, ""));

        } else if (name == "a") {
            p.in_a = true;
        }

    }

    static void endElementNs(void* ctx, const xmlChar* localname, const xmlChar* prefix, const xmlChar* URI) {

        FB2Parser& p = *(static_cast<FB2Parser*>(ctx));

        std::string name((char*)localname);

        if (name == "title-info") {
            p.in_title_info = false;

        } else if (name == "author") {
            p.in_author = false;

        } else if (name == "first-name") {
            p.in_first_name = false;

        } else if (name == "middle-name") {
            p.in_middle_name = false;

        } else if (name == "last-name") {
            p.in_last_name = false;

        } else if (name == "book-title") {
            p.in_book_title = false;

        } else if (name == "lang") {
            p.in_lang = false;

            if (p.langs.count(p.book.lang) == 0)
                p.ok = false;

        } else if (name == "body") {
            p.in_body = false;

        } else if (name == "title") {
            p.in_title = false;

        } else if (name == "subtitle") {
            p.in_subtitle = false;

        } else if (name == "p") {
            p.in_p = false;

        } else if (name == "v") {
            p.in_v = false;

        } else if (name == "text-author") {
            p.in_text_author = false;

        } else if (name == "a") {
            p.in_a = false;
        }
    }

    static void characters(void *ctx, const xmlChar* ch, int len) {

        FB2Parser& p = *(static_cast<FB2Parser*>(ctx));

        std::string chars((char*)ch, len);

        if (p.in_title_info) {

            if (p.in_author) {

                if (p.in_first_name) {
                    p.book.author.back()[0] += chars;

                } else if (p.in_middle_name) {
                    p.book.author.back()[1] += chars;

                } else if (p.in_last_name) {
                    p.book.author.back()[2] += chars;
                }

            } else if (p.in_book_title) {
                p.book.title += chars;

            } else if (p.in_lang) {
                p.book.lang += chars;
            }

        } else if (p.in_body) {

            if (!p.in_a && 
                (p.in_p || p.in_v || p.in_text_author || p.in_subtitle)) {

                p.book.text.back().second += chars;
            }

        }
    }

    static void fatalError(void* ctx, const char* msg, ...) {

        FB2Parser& fsm = *(static_cast<FB2Parser*>(ctx));
        va_list args;
        va_start(args, msg);
        vprintf( msg, args );
        va_end(args);
    }

    static void error(void* ctx, const char* msg, ...) {

        FB2Parser& fsm = *(static_cast<FB2Parser*>(ctx));
        va_list args;
        va_start(args, msg);
        vprintf( msg, args );
        va_end(args);
    }

    static void warning(void* ctx, const char* msg, ...) {

        FB2Parser& fsm = *(static_cast<FB2Parser*>(ctx));
        va_list args;
        va_start(args, msg);
        vprintf( msg, args );
        va_end(args);
    }
};

struct IdfMapper {
    static const size_t N = 1000;
    std::map<size_t,size_t> idf;

    std::string filename;

    void add(const std::string& fn, size_t n, const Book& book) {

        filename = fn;

        for (const auto& z : book.w) {
            idf[z.first] += z.second;
        }

        if ((n % N) == N - 1) {

            write_terms(idf, filename);
            idf.clear();
        }
    }

    ~IdfMapper() {

        if (filename.size() > 0) {
            write_terms(idf, filename);
        }
    }

};

IdfMapper& idf() {
    static IdfMapper _idf;
    return _idf;
}

void write_book(const Book& book) {

    size_t N = idf().N;
    auto& _idf = idf();

    size_t n = files().add(book.filename);

    std::cout << "  " << n << " " << book.title << " [" << book.lang << "] : ";
    for (const auto& z : book.author) {
        std::cout << z[0] << " " << z[1] << " " << z[2];
    }
    std::cout << std::endl;

    std::stringstream fn;
    fn << "books." << (n - (n % N)) + N;

    write_book(book, fn.str());

    fn << ".terms";

    _idf.add(fn.str(), n, book);
}

void handle_parsed(Book& book) {

    std::map<size_t,size_t> t_i;
    std::map<size_t,size_t> a_i;
    size_t i2 = 0;

    lemm()(book.title, t_i, i2);

    for (const auto& z : book.author) {

        lemm()(z[0], a_i, i2);
        lemm()(z[1], a_i, i2);
        lemm()(z[2], a_i, i2);
    }

    for (const auto& i : book.text) {

        lemm()(i.second, book.w, book.n);
    }

    for (const auto& z : t_i) {
        book.title_index.insert(z.first);
    }

    for (const auto& z : a_i) {
        book.author_index.insert(z.first);
    }

    // HACK!
    book.text.clear();

    write_book(book);
}

        
void read_zip(const std::string& zipfile, const std::set<std::string>& langs) {

    struct archive* a;
    struct archive_entry* entry;

    a = archive_read_new();
    archive_read_support_format_all(a);

    int r;

    r = archive_read_open_filename(a, zipfile.c_str(), 10240);

    if (r < 0) {
        throw std::runtime_error(std::string("Could not open archive: ") + archive_error_string(a));
    }

    while (1) {
        r = archive_read_next_header(a, &entry);

        if (r == ARCHIVE_EOF)
            break;

        if (r != ARCHIVE_OK)
            throw std::runtime_error(std::string("Failed reading archive: ") + archive_error_string(a));

        std::string filename = archive_entry_pathname(entry);

        std::cerr << filename << std::endl;

        xmlSAXHandler saxHandler;
        xmlParserCtxtPtr ctxt;

        memset(&saxHandler, 0, sizeof(saxHandler));

        saxHandler.initialized = XML_SAX2_MAGIC;

        saxHandler.startElementNs = &FB2Parser::startElementNs;
        saxHandler.endElementNs   = &FB2Parser::endElementNs;
        saxHandler.characters     = &FB2Parser::characters;
        saxHandler.warning        = &FB2Parser::warning;
        saxHandler.error          = &FB2Parser::error;
        saxHandler.fatalError     = &FB2Parser::fatalError;

        FB2Parser parser(langs);

        //int result = xmlSAXUserParseFile(&saxHandler, &parser, argv[1]);

        char buff[64*1024];
        size_t size = 64*1024;

        ctxt = xmlCreatePushParserCtxt(&saxHandler, &parser, buff, 0, filename.c_str());

        while (1) {
            r = archive_read_data(a, &buff, size);

            if (r == 0)
                break;

            if (r < 0) {
                throw std::runtime_error(std::string("Failed reading data: ") + archive_error_string(a));
            }

            r = xmlParseChunk(ctxt, buff, r, 0);

            if (r != 0) {
                std::cerr << "Bad xml: " << filename << std::endl;
            }

            if (!parser.ok)
                break;
        }

        if (parser.ok) {
            xmlParseChunk(ctxt, buff, 0, 1);
        }

        xmlFreeParserCtxt(ctxt);

        if (!parser.ok)
            continue;

        size_t n = files().has(filename);

        if (n != 0)
            continue;

        if (parser.book.text.empty())
            continue;

        parser.book.archive = zipfile;
        parser.book.filename = filename;

        handle_parsed(parser.book);
    }

    archive_read_close(a);
    archive_read_free(a);
}


int main(int argc, const char** argv) {

    //LIBXML_TEST_VERSION

    try {

        std::string zipfile;

        if (argc <= 2) {
            std::cerr << "Usage: " << argv[0] << " <language> <zipfile>..." << std::endl;
            return 1;
        }

        std::set<std::string> langs({ argv[1] });

        for (int i = 2; i < argc; ++i) {

            zipfile = argv[i];

            read_zip(zipfile, langs);
        }

    } catch (std::exception& e) {
        std::cerr << "ERROR: " << e.what() << std::endl;
        return 1;

    } catch (...) {
        std::cerr << "UNKNOWN ERROR" << std::endl;
        return 1;
    }

    return 0;
}

